Ground state simulation of the Bose-Hubbard model.

# pylint: disable=invalid-name

import numpy as np

import qtealeaves as qtl
from qtealeaves.models import get_bose_hubbard_1d

# Reference ground state energies from OSMPS
# Keys are ... L, J, symmetry_sector
ref_osmps = {
    (8, 0.5, None): -4.7358608335175,
    (8, 0.5, 8): -4.66662145271015,
    (8, 0.5, 7): -4.4656708344888,
}


def main(tn_type=5, statics_method=2, input_folder=None, output_folder=None):
    """
    Main method for the ground state calculation fo the
    1d Bose-Hubbard model.

    **Arguments**

    tn_type : int, optional
        Choose 5 for python-TTN, 6 for python-MPS.
        Default to 5.

    statics_method : integer, optional
        Method to run ground state search for this/all iteration.
        0 : default (2)
        1 : sweep
        2 : sweep with space expansion (can still be reduced to sweep
            during the simulation based on a energy condition)

    input_folder : str | None, optional
        Input folder. Default to None.

    output_folder : str | None, optional
        Output folder. Default to None.
    """
    if input_folder is None:
        input_folder = lambda params: "BOHU/input_L%d" % (params["L"])
    if output_folder is None:
        output_folder = lambda params: "BOHU/output_L%d" % (params["L"])

    model, my_ops = get_bose_hubbard_1d()

    my_conv = qtl.convergence_parameters.TNConvergenceParameters(
        max_iter=7, max_bond_dimension=16, statics_method=statics_method
    )
    my_obs = qtl.observables.TNObservables()

    simulation = qtl.QuantumGreenTeaSimulation(
        model,
        my_ops,
        my_conv,
        my_obs,
        tn_type=tn_type,
        folder_name_input=input_folder,
        folder_name_output=output_folder,
        has_log_file=False,
        store_checkpoints=False,
    )

    params = []
    params.append({"L": 8, "U": 1.0, "J": 0.5})

    simulation.run(params, delete_existing_folder=True)

    for elem in params:
        e0_tn = simulation.get_static_obs(elem)["energy"]
        e0_osmps = ref_osmps[(elem["L"], elem["J"], None)]

        print("Energies TTN vs OSMPS", e0_tn, e0_osmps)

        assert np.abs(e0_tn - e0_osmps) < 1e-4

    print(
        f"\nExample `{__file__}` ran successfully; "
        + "ground state energy at least correct up to 1e-4."
    )

    return


if __name__ == "__main__":
    main()